#include "grate_test_utils.h"
#include "dfa_construct.h"

#include "grate/dfa.h"

using growth_rate::LexminTrie;
using growth_rate::LexminDFA;
using growth_rate::NodeData;
using growth_rate::NodeQueue;
using growth_rate::RoundUpToPowerOfTwo;
using growth_rate::Log2;
using growth_rate::testing::TableDFA;

namespace {

TEST(DFAConstructTest, NodeQueueTest) {
  size_t queue_size = 10;
  NodeQueue queue(queue_size);
  EXPECT_TRUE(queue.empty());
  for (size_t i = 0; i < 12; i++) {
    EXPECT_TRUE(queue.push(i));
    EXPECT_FALSE(queue.empty());
    EXPECT_EQ(i, queue.pop());
    EXPECT_TRUE(queue.empty());
  }
  // Fill the queue completely.
  for (size_t i = 0; i < queue_size; i++) {
    EXPECT_TRUE(queue.push(i));
  }
  EXPECT_FALSE(queue.empty());
  EXPECT_FALSE(queue.push(0));
  EXPECT_EQ(0U, queue.pop());
  EXPECT_TRUE(queue.push(0));
}

static void RunNodeDataTest(size_t size, uint8_t alphabet_size) {
  TableDFA *dfa = new TableDFA(size, alphabet_size);
  NodeData *data = new NodeData(dfa);

  // Undefined letter should be distinguishable from a normal one.
  EXPECT_LE(alphabet_size, data->undefined_letter());
  // Check that data size is small enough.
  size_t exp_letter_size = Log2(RoundUpToPowerOfTwo(alphabet_size));
  size_t expected_bit_size = Log2(RoundUpToPowerOfTwo(size)) +
                             exp_letter_size +
                             alphabet_size * (1 + exp_letter_size);
  size_t actual_bit_size = data->node_size_in_bits();
  EXPECT_LE(actual_bit_size, expected_bit_size + 10);
  size_t expected_total_size = size * ((actual_bit_size + 7) / 8);
  size_t actual_total_size = data->data_size_in_bytes();
  EXPECT_LE(actual_total_size, (size_t)(1.1 * expected_total_size));

  // Check setters/getters.
  for (size_t node = 0; node < size; node++) {
    size_t lf = random(node + 1);
    data->set_lf(node, lf);
    EXPECT_EQ(lf, data->get_lf(node));
    uint8_t last = random(alphabet_size);
    data->set_last(node, last);
    EXPECT_EQ(last, data->get_last(node));
    for (uint8_t letter = 0; letter < alphabet_size; letter++) {
      EXPECT_EQ(data->undefined_letter(), data->get_sigma(node, letter));
      uint8_t sigma = random(alphabet_size);
      data->set_sigma(node, letter, sigma);
      EXPECT_EQ(sigma, data->get_sigma(node, letter));
      bool back = random(2);
      data->set_back(node, letter, back);
      EXPECT_EQ(back, data->get_back(node, letter));
    }
  }

  // Test auxiliary sigma methods.
  size_t src_node = random(size);
  size_t dest_node = random(size);
  uint8_t sigma_array[256];
  uint8_t sigma2[256];
  data->copy_sigma(dest_node, src_node);
  data->copy_sigma_to_array(src_node, sigma_array);
  for (uint8_t letter = 0; letter < alphabet_size; letter++) {
    EXPECT_EQ(data->get_sigma(src_node, letter), sigma_array[letter]);
    EXPECT_EQ(data->get_sigma(dest_node, letter), sigma_array[letter]);
    sigma2[letter] = random(alphabet_size);
  }
  data->apply_to_sigma(src_node, sigma2);
  for (uint8_t letter = 0; letter < alphabet_size; letter++) {
    if (sigma_array[letter] < alphabet_size) {
      EXPECT_EQ(sigma2[sigma_array[letter]], data->get_sigma(src_node, letter));
    }
  }

  delete data;
  delete dfa;
}

TEST(DFAConstructTest, NodeDataTest) {
  RunNodeDataTest(10, 7);
  RunNodeDataTest(5, 2);
  RunNodeDataTest(20, 15);
  RunNodeDataTest(10000, 150);
  RunNodeDataTest(100000, 3);
}

TEST(DFAConstructTest, NodeDataDumpTest) {
  TableDFA *dfa = new TableDFA(10, 3);
  NodeData *data = new NodeData(dfa);

  size_t node = 3;
  data->set_lf(node, 2);
  data->set_last(node, 1);
  data->set_sigma(node, 0, 0);
  data->set_sigma(node, 1, 1);
  data->set_sigma(node, 2, 2);
  data->set_back(node, 0, true);
  data->set_back(node, 1, false);
  data->set_back(node, 2, true);
  EXPECT_EQ("[node = 3: lf = 2, last = 1, sigma = <0, 1, 2>, "
            "back = <1, 0, 1>]", data->dump_node(node));

  delete data;
  delete dfa;
}

TEST(DFAConstructTest, Construction) {
  {
    LexminTrie chain_trie(10, 2);
    EXPECT_TRUE(chain_trie.add_lexmin_word("aaaa"));
    LexminDFA chain_dfa(&chain_trie);
    int exp_table[] = {1, 1,
                       2, 1,
                       3, 1,
                       -1, 1};
    TableDFA exp_dfa(4, 2, exp_table);
    EXPECT_TRUE(exp_dfa.equals(chain_dfa));
  }

  {
    LexminTrie small_trie(20, 3);
    EXPECT_TRUE(small_trie.add_lexmin_word("abaab"));
    EXPECT_TRUE(small_trie.add_lexmin_word("aabc"));
    EXPECT_TRUE(small_trie.add_lexmin_word("abbc"));
    LexminDFA small_dfa(&small_trie);
    int exp_table[] = {1, 1, 1,
                       5, 2, 2,
                       3, 7, 2,
                       4, 3, 2,
                       5, -1, -1,
                       5, 6, 6,
                       3, 7, -1,
                       6, 5, -1};
    TableDFA exp_dfa(8, 3, exp_table);
    EXPECT_TRUE(exp_dfa.equals(small_dfa));
  }

  {
    LexminTrie big_alph_trie(10, 7);
    EXPECT_TRUE(big_alph_trie.add_lexmin_word("ab"));
    LexminDFA big_alph_dfa(&big_alph_trie);
    int exp_table[] = {1, 1, 1, 1, 1, 1, 1,
                       1, -1, -1, -1, -1, -1, -1};
    TableDFA exp_dfa(2, 7, exp_table);
    EXPECT_TRUE(exp_dfa.equals(big_alph_dfa));
  }

  {
    LexminTrie trie(10, 2);
    EXPECT_TRUE(trie.add_lexmin_word("abb"));
    EXPECT_TRUE(trie.add_lexmin_word("aabb"));
    LexminDFA dfa(&trie);
    int exp_table[] = {1, 1,
                       3, 2,
                       2, -1,
                       3, 4,
                       2, -1};
    TableDFA exp_dfa(5, 2, exp_table);
    EXPECT_TRUE(exp_dfa.equals(dfa));
  }
}

}  // namespace
